from io import TextIOWrapper
import os
import re
import fcntl
import os.path as osp
import subprocess
from typing import Dict, List, Optional
import argparse
import sys
import pwd

__version__ = '0.0.7'

DEVICE_LOCK_ROOT_ENV = 'DEVICE_LOCKS'

LOCK_DIR = os.environ.get(DEVICE_LOCK_ROOT_ENV)

def is_lock_enabled():
    '''
    return whether device lock system can be use.
    you must set `DEVICE_LOCKS` environment variable to use this device lock system.
    '''
    return LOCK_DIR is not None

__global_locks: Dict[str, TextIOWrapper] = dict()

def gen_lock_files(nums=None):
    '''
    generate lock files.
    nums: `int | None`. device count, if None, will work out by nvidia-smi.
    '''
    if not is_lock_enabled():
        raise ValueError(f"env variable f{DEVICE_LOCK_ROOT_ENV} not set, can't gen lock files")

    if nums is not None:
        if type(nums) is not int:
            raise TypeError("nums must be int or None")
        if LOCK_DIR is not None:
            for i in range(nums):
                file_name = osp.join(LOCK_DIR, str(i))
                with open(file_name, 'w') as f:
                    f.write(str(i))
    else:
        out = subprocess.run(['nvidia-smi', '-L'], check=True, capture_output=True)
        lines = out.stdout.strip().split(b'\n')
        nums = len(lines)
        for i in range(nums):
            file_name = osp.join(LOCK_DIR, str(i))
            with open(file_name, 'w') as f:
                f.write(str(i))

def _get_device_index() -> List[int]:
    if not is_lock_enabled():
        return []
    devices = [int(f) for f in os.listdir(LOCK_DIR)]
    devices.sort()
    return devices

def _index2lock(index: int) -> str:
    return osp.join(LOCK_DIR, str(index))

def _get_device_file() -> List[str]:
    if not is_lock_enabled():
        return []
    devices = _get_device_index()
    files = [_index2lock(d) for d in devices]
    return files

def _get_file_inode(file_path: str) -> int:
    out = os.stat(file_path)
    return out.st_ino

def get_device_count() -> int:
    '''
    return all device under control.
    '''
    if not is_lock_enabled():
        return 0
    return len(_get_device_file())

def _fetch_lock(file_path: str) -> Optional[TextIOWrapper]:
    file = open(file_path, 'r')
    try:
        fcntl.flock(file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        return None
    return file

def _lock_device(index: int) -> Optional[TextIOWrapper]:
    if not is_lock_enabled():
        return None
    lock_file_path = osp.join(LOCK_DIR, str(index))
    return _fetch_lock(lock_file_path)


def scan_locks() -> Dict[str, str]:
    '''
    return all devices and occupied process.
    '''
    if not is_lock_enabled():
        return None

    inode_map = {_get_file_inode(f) : i for i, f in enumerate(_get_device_file())}
    ret = dict()
    with open('/proc/locks', 'r') as f:
        lines = f.readlines()
    lines = [l.split() for l in lines]
    for line in lines:
        pid = int(line[4])
        file_inode = int(line[5].split(':')[2])
        if file_inode in inode_map:
            device = inode_map[file_inode]
            if device in ret:
                ret[device].append({"pid", pid})
            else:
                ret[device] = [{"pid": pid}]
    return ret

def mask_device(indexes: List[int]):
    '''
    must mask device in `CUDA_VISIBLE_DEVICES`.
    '''
    out = indexes.copy()
    out.sort()
    ENV = ','.join([str(i) for i in out])
    os.environ['CUDA_VISIBLE_DEVICES'] = ENV

def fetch_device(count: int, allow_shrink=False, mask=False) -> Optional[List[int]]:
    '''
    try to lock some device for use.
    count: int. max device count required. 0 if require all.
    allow_shrink=False: bool. allow return less device than required.
    mask=False: bool. also set `CUDA_VISIBLE_DEVICES` environment variable for locked device.
    
    return: return None if failed, else return actual locked device number.
    '''
    if not is_lock_enabled():
        return None
    if count < 0:
        raise ValueError("count can't be negative")
    files = _get_device_file()
    if count > len(files):
        raise ValueError(f"require too many devices, total {len(files)}, require {count}")

    locks = []
    devices = []
    for i in _get_device_index():
        if len(locks) == count and count != 0:
            break
        lock = _lock_device(i)
        if lock is not None:
            locks.append(lock)
            devices.append(i)

    if len(locks) < count and not allow_shrink:
        for l in locks:
            l.close()
        return None

    for d, l in zip(devices, locks):
        __global_locks[d] = l
    if mask:
        mask_device(devices)
    return devices

def release_all() -> List[int]:
    '''
    release all locked device. will not recovery `CUDA_VISIBLE_DEVICES` environment variable.
    your process will automatically release device lock on exit.
    you don't need to call this function explicitly.
    '''
    ret = []
    for d, l in __global_locks:
        l.close()
        ret.append(d)
    __global_locks = dict()
    return ret

def _find_process_owner(pid: int) -> Optional[str]:
    with open(f'/proc/{pid}/status', 'r') as f:
        lines = f.readlines()
    for line in lines:
        if not line.startswith('Uid:'):
            continue
        uid = int(line.split()[1])
        out = pwd.getpwuid(uid)
        return out.pw_name
    return None

if __name__ == '__main__':

    def handle_gen(args):
        gen_lock_files(nums=args.c)
    
    DOCKER_CGROUP_PATTERN = re.compile(r'docker-(.*).scope')
    def check_docker_process(cgroups: list[str]) -> Optional[str]:
        for line in cgroups:
            out = DOCKER_CGROUP_PATTERN.search(line)
            if out is not None:
                container_id = out.group(1)
                try:
                    name_res = subprocess.run([
                        'docker', 'inspect', '--format', '{{.Name}}', container_id
                    ], check=True, capture_output=True)
                    container_name = name_res.stdout.strip().decode('utf-8')[1:]
                except Exception:
                    container_name = "ERROR"
                return f"docker::{container_name}::{container_id}"
        return None

    def find_process_container_id(pid: int) -> Optional[str]:
        with open(f'/proc/{pid}/cgroup', 'r') as f:
            lines = f.readlines()
        return check_docker_process(lines)


    def show_occpy(d: int, ps: List[Dict]):
        # sort process by pid
        ps.sort(key=lambda x: x['pid'])
        for p in ps:
            try:
                user_name = _find_process_owner(p['pid'])
            except Exception:
                user_name = "ERROR"

            try:
                out = find_process_container_id(p['pid'])
            except Exception:
                out = None

            if out is None:
                print(f"device #{d} :: PID: {p['pid']} USER: {user_name}")
            else:
                print(f"device #{d} :: PID: {p['pid']} USER: {user_name} {out} ")


    def handle_list(args):
        out = list(scan_locks().items())

        # sort device by device id
        out.sort(key=lambda x: x[0])
        devices = args.d
        if devices is None:
            for k, v in out:
                show_occpy(k, v)
        else:
            for k, v in out:
                if k in devices:
                    show_occpy(k, v)

    parser = argparse.ArgumentParser(prog='device_lock')
    parser.set_defaults(func=lambda _: parser.print_help(sys.stdout))
    subparsers = parser.add_subparsers(title='subcommands', description='valid subcommands', dest='cmd')
    gen_parser = subparsers.add_parser('gen', help='gen command gen lock files')
    gen_parser.add_argument('-c', type=int, default=None, help='direct set device count')
    gen_parser.set_defaults(func=handle_gen)
    list_parser = subparsers.add_parser('list', help="list all occupied device related process, don't use in container")
    list_parser.add_argument('-d', type=int, nargs='+', default=None, help='device to list')
    list_parser.set_defaults(func=handle_list)
    ARGS = parser.parse_args()
    ARGS.func(ARGS)